package algotithm.leetcode.tree.test814;

import algotithm.leetcode.tree.TreeNode;

/**
 * @author zhouyanxiang
 * @Date 2021-02-2021/2/25-14:26
 * @URL https://leetcode-cn.com/problems/binary-tree-pruning/
 * @Title 814. 二叉树剪枝
 */
public class Solution {

    public TreeNode pruneTree(TreeNode root) {
        return containsOne(root) ? root : null;
    }

    public boolean containsOne(TreeNode root) {
        if (root == null) {
            return false;
        }
        boolean a1 = containsOne(root.left);
        boolean a2 = containsOne(root.right);
        if (!a1) {
            root.left = null;
        }
        if (!a2) {
            root.right = null;
        }
        return root.val == 1 || a1 || a2;
    }
}
